# -*- coding:utf-8 -*-
import os,sys
import re
import traceback
import time
import copy
from collections import deque
sys.path.append(os.path.join(os.path.abspath(os.path.dirname(__file__)), os.pardir))
import supeanut_config
from CommonLib.mylog import mylog
from CommonLib.StockCal.BasicTool import BasicTool
from CommonLib.StockCal.ExrightsTool import ExrightsTool
from CommonLib.db.MongoDBTool import MongoDBTool
from CommonLib.StockTool import StockTool
from CommonLib.StockCal.TecIndexRecog import TecIndexRecog

'''
作者：supeanut
创建时间：2017-06-20 xx:xx:xx
功能描述：
相关配置：
	supeanut_config.XXX
历史改动：
	2017-xx-xx: xxxxxx
'''
class TrainDataSupport:
	def __init__(self):
		self.hq_obj = BasicTool()
		self.adj_obj = ExrightsTool()
		self.db_obj = MongoDBTool()
		self.tec_obj = TecIndexRecog()

	# 区间前复权
	# exData_list 必须按照 时间先后
	def preCal(self, ori_change, itemList, exData_list):
		if ori_change is False:
			adj_itemList = copy.deepcopy(itemList)
		else:
			adj_itemList = itemList
		for exrights_data in exData_list:
			paixi = exrights_data["paixi"]
			peigu = exrights_data["peigu"]
			peigujia = exrights_data["peigujia"]
			songgu = exrights_data["songgu"]
			zenggu = exrights_data["zenggu"]
			for item in adj_itemList:
				for i in [1,2,3,4]:
					item[i] = (item[i] - paixi + peigu * peigujia) / ( 1.0 + songgu + zenggu + peigu)
				item[5] = item[5] * (1.0 + songgu + zenggu + peigu)
		return adj_itemList

	# 次新股连涨过滤
	def filter_upComtinue(self, itemList):
		lenItemList = len(itemList)
		# 不分析仅存在一天的股票
		if lenItemList < 2:
			return []
		close_pre = itemList[0][2]
		for i in range(1, lenItemList):
			item = itemList[i]
			close = item[2]
			if (close - close_pre)/close_pre < 0.098:
				break
			close_pre = close
		# 至今未开板
		if lenItemList == (i+1):
			return []
		return itemList[i:]

	# 历史行情分段
	# return {date:{'itemList':itemList, 'ex_info':ex_info}, ...}
	def getOriHQ(self, code, index_tag, period_s, period_e):
		# 获取原始行情和除权数据
		flag, itemList = self.hq_obj.getOriItemList(code, index_tag, "day", period_s, period_e)
		if flag is False:
			return False, itemList
		# 次新股连涨过滤
		itemList = self.filter_upComtinue(itemList)
		flag, exData = self.adj_obj.getExrights(code)
		if flag is False:
			return False, exData
		if len(itemList) == 0:
			return True, {}
		ex_dates = exData.keys()
		ex_dates.sort()
		len_ex_date = len(ex_dates)
		for item in itemList:
			item[0] = item[0][:10]
		if len_ex_date == 0:
			return True, {itemList[-1][0]:{'itemList':itemList,'ex_info':None}}
		# 第一个行情日 逻辑上的preDate
		first_ori_date = itemList[0][0][:10]
		s_ex_date = ex_dates[0]
		oriPreDate = '0000-00-00'
		if s_ex_date < first_ori_date:
			oriPreDate = s_ex_date
		# 按照除权日分段
		itemDict_byPeriod = {}
		itemList_byPeriod = []
		# --定位第一个s_ex_date
		s_ex_date_index = 0
		for i in range(0, len_ex_date):
			if ex_dates[i] >= first_ori_date:
				s_ex_date_index = i
				break
		s_ex_date = ex_dates[s_ex_date_index]
		for item in itemList:
			date = item[0]
			# 判断当前日是否除权日,考虑非交易情况
			if s_ex_date > oriPreDate and s_ex_date <= date:
				# 保存前一个除权区间,新建最新的除权区间
				if len(itemList_byPeriod) > 0:
					itemDict_byPeriod[oriPreDate] = {'itemList':itemList_byPeriod,'ex_info':exData.get(s_ex_date)}
				itemList_byPeriod = [item]
				s_ex_date_index += 1
				if s_ex_date_index < len_ex_date:
					s_ex_date = ex_dates[s_ex_date_index]
				else:
					s_ex_date = '9999-99-99'
			else:
				itemList_byPeriod.append(item)
			oriPreDate = date
		if len(itemList_byPeriod) > 0:
			itemDict_byPeriod[oriPreDate] = {'itemList':itemList_byPeriod,'ex_info':None}
		return True, itemDict_byPeriod

	# 历史行情还原(前复权)
	def recPreHQ(self, itemDict_byPeriod, ):
		itemDict_byIter = {}
		e_dates = itemDict_byPeriod.keys()
		e_dates.sort()
		s_dates = []
		itemLists = []
		ex_infos = []
		date_periods_str = []
		for e_date in e_dates:
			itemDict = itemDict_byPeriod[e_date]
			itemList = itemDict['itemList']
			ex_info = itemDict['ex_info']
			s_date = itemList[0][0][:10]
			itemLists.append(itemList)
			s_dates.append(s_date)
			ex_infos.append(ex_info)
			date_periods_str.append('%s_%s'%(s_date,e_date))
		# 创建原始itemDict_byIter
		len_itemDict = len(e_dates)
		for i in range(len_itemDict):
			date_period_str = date_periods_str[i]
			itemDict_byIter[date_period_str] = [k for k in copy.deepcopy(itemLists[:i+1])]
		# 迭代还原行情
		for j in range(0, len_itemDict-1):
			for i in range(j+1, len_itemDict):
				for j_iter in range(0, j+1):
					itemList = itemDict_byIter[date_periods_str[i]][j_iter]
					self.preCal(True, itemList, [ex_infos[j]])
		# 合并list of itemList
		for i in range(len_itemDict):
			itemLists = itemDict_byIter[date_periods_str[i]]
			all_items = []
			for itemList in itemLists:
				all_items += itemList
			itemDict_byIter[date_periods_str[i]] = all_items
		return itemDict_byIter

	# 计算label，返回{date:label}
	def calLabel(self, itemDict_byIter, indexItemDict_byIter, label_fun_param_list, stop_dates):
		date_label_dicts = []
		date_labels_dict = {}
		# 计算label
		for label_fun_param in label_fun_param_list:
			label_name = label_fun_param[0]
			label_param = label_fun_param[1]
			date_label_dict = getattr(self, label_name)(itemDict_byIter,indexItemDict_byIter,label_param, stop_dates)
			date_label_dicts.append(date_label_dict)
		# 合并label
		if len(date_label_dicts) > 1:
			jiaoji_date = set(date_label_dicts[0].keys())
			for date_label_dict in date_label_dicts[1:]:
				dates = set(date_label_dict.keys())
				jiaoji_date = jiaoji_date & dates
			jiaoji_date_list = list(jiaoji_date)
		else:
			jiaoji_date_list = date_label_dicts[0].keys()
		for common_date in jiaoji_date_list:
			date_labels_dict[common_date] = [k[common_date] for k in date_label_dicts]
		return date_labels_dict

	# 相对指数收益
	def RelaProfit(self, itemDict_byIter, indexItemDict_byIter, label_param, stop_dates):
		label_dict = self.Profit(itemDict_byIter, indexItemDict_byIter, label_param, stop_dates)
		# index无stop_dates ！
		index_label_dict = self.Profit(indexItemDict_byIter, indexItemDict_byIter, label_param, [])
		rela_profit_dict = {}
		for date in label_dict:
			# 判断指标数据比个股数据短的日期缺失
			index_profit = index_label_dict.get(date)
			if index_profit is None:
				continue
			# 个股持股停盘：holdProfit-CompleteIndexProfit；个股未买停盘：0.0-CompleteIndexProfit
			rela_profit_dict[date] = (label_dict[date] - index_label_dict[date])
		return rela_profit_dict
		
		
	# 指定日第二天开盘买，持股第days天收盘卖
	def Profit(self, itemDict_byIter, indexItemDict_byIter, label_param, stop_dates):
		label_dict = {}
		days = label_param.get('days')
		periods = itemDict_byIter.keys()
		periods.sort()
		# 滑动小窗口
		period_s_index = 0
		period_e_index = -1
		period_i = 0
		for period in periods:
			itemList_iter = itemDict_byIter[period]
			len_itemList = len(itemList_iter) - (period_e_index+1)
			period_e_index = period_s_index + len_itemList - 1
			# 除权位置处理
			if period_i == 0:
				for cur_i in range(period_s_index, period_e_index-days+1):
					date = itemList_iter[cur_i][0]
					profit_value = self.Profit_window(itemList_iter[(cur_i+1):(cur_i+1+days)],date,stop_dates)
					label_dict[date] = profit_value
			else:
				if (period_s_index-days) < 0:
					cur_i_s = 0
				else:
					cur_i_s = period_s_index-days
				for cur_i in range(cur_i_s, period_e_index-days+1):
					date = itemList_iter[cur_i][0]
					profit_value = self.Profit_window(itemList_iter[(cur_i+1):(cur_i+1+days)],date,stop_dates)
					label_dict[date] = profit_value
			# 循环尾部处理
			period_i += 1
			period_s_index += len_itemList
		return label_dict

	# 小窗口，开买收卖
	def Profit_window(self, itemWindow, select_date, stop_dates):
		buy_item = itemWindow[0]
		sell_item = itemWindow[-1]
		buy_price = buy_item[1]
		sell_price = sell_item[2]
		buy_date = buy_item[0]
		sell_date = sell_item[0]
		for stop_date in stop_dates:
			if stop_date > select_date and stop_date < buy_date:
				# 没有机会买入，所以返回0.0中性结果
				return 0.0
			if stop_date > buy_date and stop_date < sell_date:
				# 持股到一半，结算盈利
				hold_date_e_index = 0
				for i in range(0, len(itemWindow)):
					if itemWindow[i][0] > stop_date:
						hold_date_e_index = i - 1
						break
				sell_price = itemWindow[hold_date_e_index][2]
				return (sell_price - buy_price) / buy_price
		# 中途无停牌的情况
		return (sell_price - buy_price) / buy_price
			

	# 返回 '2011-01-01_2011-09-09':all_from_start[tec1_list, tec2_list, ...]
	def calTec(self, itemDict_byIter, tec_info_list):
		tecDict_byIter = {}
		for period in itemDict_byIter:
			tecDict_byIter[period] = []
			itemList = itemDict_byIter[period]
			for tec_info in tec_info_list:
				tec_name = tec_info[0]
				tec_param = tec_info[1]
				tec_lists = getattr(self.tec_obj, tec_name)(itemList, **tec_param)
				# 判断多返回的技术指标
				if len(tec_lists) == 0:
					tecDict_byIter[period].append(tec_lists)
				elif (type(tec_lists[0]) == list):
					for tec_list in tec_lists:
						tecDict_byIter[period].append(tec_list)
				else:
					tecDict_byIter[period].append(tec_lists)
		return tecDict_byIter
			
	# 返回 {date :[shape1,shape2,], }
	def calShape(self, itemDict_byIter, tecDict_byIter, shape_info_list, used_period_index_inFile):
		date_shapes_dict = {}
		# 获取最大windowSize
		bigWinSize = 0
		for shape_info in shape_info_list:
			if bigWinSize < shape_info[2]:
				bigWinSize = shape_info[2]
		# 按时间遍历ItemIter
		periods = itemDict_byIter.keys()
		periods.sort()
		period_s_i = 0
		for period_str in periods:
			itemList = itemDict_byIter[period_str]
			tecLists = tecDict_byIter[period_str]
			itemList_len = len(itemList)
			# 计算shapes
			for shape_info in shape_info_list:
				shape_name = shape_info[0]
				tecs_index = shape_info[1]
				window_size = shape_info[2]
				shape_param = shape_info[3]
				for win_e in range(period_s_i, itemList_len):
					win_s = win_e - bigWinSize + 1
					if win_s < 0:
						continue
					# 计算小窗口window下的shape
					# itemWindows  [[tec1,tec2,[open,close,high,low,vol]], ...]
					itemWindows = []
					for cur_i in range(win_s, win_e+1):
						itemWindow = []
						# 加载每日tec
						for tec_index in tecs_index:
							itemWindow.append(tecLists[tec_index][cur_i])
						# 加载每日HQ
						itemWindow.append(itemList[cur_i][1:])
						itemWindows.append(itemWindow)
					# 返回的shapes是list
					shapes = getattr(self.tec_obj, shape_name)(itemWindows, window_size, **shape_param)
					e_date = itemList[win_e][0]
					if date_shapes_dict.has_key(e_date):
						date_shapes_dict[e_date].extend(shapes)
					else:
						date_shapes_dict[e_date] = shapes
			period_s_i = itemList_len
		return self.filter_tecPreHot(date_shapes_dict, used_period_index_inFile)
	
	# 过滤技术指标平滑预热区域,如MA610
	def filter_tecPreHot(self, date_shapes_dict, used_period_index_inFile):
		dates = date_shapes_dict.keys()
		dates.sort()
		new_dict = {}
		i = 0
		for date in dates:
			i += 1
			if i < used_period_index_inFile:
				continue
			new_dict[date] = date_shapes_dict[date]
		return new_dict


	def storeTrainData(self, trainData_filedir, code, date_shapes_dict, date_labels_dict):
		try:
			os.makedirs(trainData_filedir)
		except:
			pass
		dates = date_shapes_dict.keys()
		dates.sort()
		file_h = open(trainData_filedir+code,'w')
		# 写入特征个数、类标个数，作为第一行
		has_common = False
		for date in dates:
			label_list = date_labels_dict.get(date)
			if label_list is None:
				continue
			has_common = True
			shape_list = date_shapes_dict[date]
			label_len = len(label_list)
			shape_len = len(shape_list)
			break
		# 没有搭配成功
		if has_common is False:
			return 
		file_h.write('%d,%d\n'%(shape_len, label_len))
		# 写入记录
		str_buf = ""
		for date in dates:
			label_list = date_labels_dict.get(date)
			if label_list is None:
				continue
			shape_list = date_shapes_dict[date]
			# 验证长度
			if len(shape_list) <> shape_len or len(label_list) <> label_len:
				raise Exception("%s,%s  shape len or label len error"%(code,date))
			str_buf += (date+','+','.join([str(k) for k in shape_list])+','+','.join([str(k) for k in label_list])+'\n')
		file_h.write(str_buf)

	def loadTrainData(self, trainData_filedir, code):
		if not os.path.isfile(trainData_filedir + code):
			return False
		file_h = open(trainData_filedir + code, 'r')
		line_i = 0
		lines = file_h.readlines()
		len_lines = len(lines)
		if len_lines == 0:
			return []
		file_head = lines[0].split(',')
		shape_num = int(file_head[0])
		label_num = int(file_head[1])
		shape_end_index = shape_num + 1
		records = []
		for line in lines[1:]:
			if line == '\n':
				continue
			line_split = line.split(',')
			date = line_split[0]
			shape_list = [int(k) for k in line_split[1:shape_end_index]]
			label_list = [float(k) for k in line_split[shape_end_index:]]
			records.append([date, shape_list, label_list])
		return records

			

if __name__ == '__main__':
	stock_obj = StockTool()
	flag, codes = stock_obj.get_all_stock_code('all')
	obj = TrainDataSupport()

	# 初始化分析参数
	codes = ['000099']
	trainData_filedir = r'/mnt/PredictModel/PeanutNo1/temp/traintemp/'
	period_s = '2000-06-22'
	period_e = '2017-06-22'
	tec_param_list = [['MACD',{'shortN':12,'longN':26,'midN':9}],['MA',{'ma_len':5}],['MA',{'ma_len':10}],['MA',{'ma_len':20}],['MA',{'ma_len':610}]]
	used_period_index_inFile = 610
	shape_param_list = [['MACDShapeV2',[0,1,2],3,{}],['MAShapeV2',[3,4,5,6],3,{}]]
	label_param_list = [['Profit', {'days':3}],['Profit', {'days':4}],['RelaProfit',{'days':3}]]
	# 相关指标格式化计算
	relaIndexCode = '000300'
	flag, indexItemDict_byPeriod = obj.getOriHQ(relaIndexCode, True, period_s, period_e)
	indexItemDict_byIter = obj.recPreHQ(indexItemDict_byPeriod)
	# 停牌信息计算
	code_stopDates_dict = {}
	for code in codes:
		flag, code_stopDates_dict[code] = stock_obj.getStopDates(code, )
		if flag is False:
			print 'get stop info error'
			exit(1) 
	# 个股 (特征,类标) 计算
	count = 0
	for code in codes:
		count += 1
		print 'prosses %s[%d]'%(code,count)
		# 生成训练数据
		flag, itemDict_byPeriod = obj.getOriHQ(code, False, period_s, period_e)
		itemDict_byIter = obj.recPreHQ(itemDict_byPeriod)
		tecDict_byIter = obj.calTec(itemDict_byIter, tec_param_list)
		date_shapes_dict = obj.calShape(itemDict_byIter, tecDict_byIter, shape_param_list, used_period_index_inFile)
		date_labels_dict = obj.calLabel(itemDict_byIter, indexItemDict_byIter, label_param_list, code_stopDates_dict[code])
		obj.storeTrainData(trainData_filedir, code, date_shapes_dict, date_labels_dict)
		# 测试用文件加载
		#trainData = obj.loadTrainData(trainData_filedir, code)
		'''
		# 测试用打印
		dates = date_shapes_dict.keys()
		dates.sort()
		for date in dates:
			labels = date_labels_dict.get(date)
			if labels is None:
				continue
			print date,'labels',date_labels_dict[date],'features',date_shapes_dict[date]
		'''
